#!/usr/bin/env python
# Author: veelion

"""
URL Pool for crawler to manage URLs
"""

import pickle
import leveldb
import time
import urllib.parse as urlparse


RED = '\x1b[31m'
GRE = '\x1b[32m'
BRO = '\x1b[33m'
BLU = '\x1b[34m'
PUR = '\x1b[35m'
CYA = '\x1b[36m'
WHI = '\x1b[37m'
NOR = '\x1b[0m'


class UrlDB:
    '''Use LevelDB to store URLs what have been done(succeed or faile)
    3个功能：存储成功的url{url:1},
            存储失败的url{url:0},
            查询准备抓取的url是否有已经存在leveldb中。
    '''
    status_failure = b'0' #请求失败的url结构：{url:0}
    status_success = b'1' #请求失败的url结构：{url:1}

    def __init__(self, db_name): #传入数据库名称
        self.name = db_name + '.urldb'
        self.db = leveldb.LevelDB(self.name) #新建leveldb数据库，如果有则进行增量存入

    def set_success(self, url):
        '''存储成功的url{url:1}'''
        if isinstance(url, str): #判断url类型
            url = url.encode('utf8') #转utf8编码
        try:
            self.db.Put(url, self.status_success) #url存入leveldb中
            s = True
        except:
            s = False
        return s

    def set_failure(self, url):
        '''存储失败的url{url:0},'''
        if isinstance(url, str):
            url = url.encode('utf8')
        try:
            self.db.Put(url, self.status_failure)
            s = True
        except:
            s = False
        return s

    def has(self, url):
        '''查询准备抓取的url是否有已经存在leveldb中。'''
        if isinstance(url, str):
            url = url.encode('utf8')
        try:
            attr = self.db.Get(url)
            return attr
        except:
            pass
        return False


class UrlPool:
    '''URL Pool for crawler to manage URLs
    '''

    def __init__(self, pool_name):
        self.name = pool_name
        self.db = UrlDB(pool_name) #调用UrlDB类，并传入leveldb数据库名称

        self.waiting = {}  # {host: set([urls]), } 按host分组，记录等待下载的URL，set([urls]):set把url列表去重
        self.pending = {}  # {url: pended_time, } 记录已被取出（self.pop()）但还未被更新状态（正在下载）的URL, pended_time是request时的时间戳
        self.failure = {}  # {url: times,} 记录失败的URL与失败次数
        self.failure_threshold = 3 #设置失败的极限次数为3次
        self.pending_threshold = 10  # 设置pending的最大时间为10秒，过期要重新下载
        self.waiting_count = 0  # self.waiting 字典里面的url的个数
        self.max_hosts = ['', 0]  # [host: url_count] 目前waiting中url最多的host及其url数量
        self.hub_pool = {}  # {url: last_query_time}  存放hub url，last_query_time：request时的时间戳，初始值为0
        self.hub_refresh_span = 0 #设置网站首页抓取间隔时间，默认值为0
        self.load_cache() #预加载，继续上一次的爬虫任务，实现断点续爬：把上次未完成的url池从硬盘加载到内存

    def __del__(self):
        '''断点续爬之：正常关闭爬虫时自动把把未完成的url池从内存保存到硬盘，实现持久化url池，为下次开启时预加载数据'''
        self.dump_cache()

    def load_cache(self,):
        '''#预加载，继续上一次的爬虫任务，实现断点续爬：把上次未完成的url池从硬盘加载到内存，'''
        path = self.name + '.pkl'
        try:
            with open(path, 'rb') as f:
                self.waiting = pickle.load(f)
            cc = [len(v) for k, v in self.waiting]
            print('saved pool loaded! urls:', sum(cc))
        except:
            pass

    def dump_cache(self):
        '''正常关闭爬虫时自动把把未完成的url池从内存保存到硬盘'''
        path = self.name + '.pkl'
        try:
            with open(path, 'wb') as f:
                pickle.dump(self.waiting, f)
            print('self.waiting saved!')
        except:
            pass

    def set_hubs(self, urls, hub_refresh_span):
        '''处理首页url列表，并设置爬取间隔时间'''
        self.hub_refresh_span = hub_refresh_span #接入网站首页抓取间隔时间
        self.hub_pool = {} #网站首页url池，{url: last_query_time},last_query_time：request时的时间戳，初始值为0
        for url in urls:
            self.hub_pool[url] = 0 #last_query_time初始值为0

    def set_status(self, url, status_code):
        '''处理爬虫返回的url与http状态码'''
        if url in self.pending: #如果url在准备爬取的url池中
            self.pending.pop(url) #则从pending的url池中取出，简单的说：就是从正在请求中的url池中删除

        if status_code == 200: #如果请求成功，按成功的url存入leveldb中
            self.db.set_success(url) #按住ctrl不放，左击set_success，可查看函数组成
            return
        if status_code == 404: #如果请求的网页不存在，按失败的url存入leveldb中
            self.db.set_failure(url)
            return
        if url in self.failure: #如果url在失败的url池中
            self.failure[url] += 1 #失败次数+1
            if self.failure[url] > self.failure_threshold: #如果失败次数超过阀值，按失败的url存入leveldb中
                self.db.set_failure(url)
                self.failure.pop(url) #按失败的url存入leveldb中后，并从失败url池中删除，pop就是删除的意思
            else:
                self.add(url) #如果url失败的url池中,并且失败次数没有超过阀值，则重新加入url池中，准备再次给爬虫爬取
        else:
            self.failure[url] = 1 #如果url不在失败的url池中，则加入失败的url池中，并记数1次
            self.add(url) #重新加入url池中，准备再次给爬虫爬取


    def push_to_pool(self, url):
        '''处理html中获取的url并往url池中加入url'''
        host = urlparse.urlparse(url).netloc #解析url并获取主机（host)
        if not host or '.' not in host: #如果url没用主机，或者为非法主机也就是没有.
            print('try to push_to_pool with bad url:', url, ', len of ur:', len(url))
            return False
        if host in self.waiting: #如果这个url的主机在url池中 {host:set(url列表)}
            if url in self.waiting[host]: #如果url已经存在url池中，则返回True，也就是不处理了
                return True
            self.waiting[host].add(url) #如果url没有经存在url池中，则加入url池中，按host分组加入{host:set(url列表)}
            if len(self.waiting[host]) > self.max_hosts[1]: #获取url数量最多的hosts
                self.max_hosts[1] = len(self.waiting[host])
                self.max_hosts[0] = host
        else: #如果这个url的主机不在url池中
            self.waiting[host] = set([url])  # host: set([urls]),创建新的键值
        self.waiting_count += 1 # self.waiting 字典里面的url的个数+1
        return True

    def add(self, url, always=False):
        '''判断url是否请求超时是否在leveldb中，并往url池中添加'''
        if always: #如果always为Ttue, 刚直接把url加入到url池中
            return self.push_to_pool(url)
        pended_time = self.pending.get(url, 0) #获取url在pending池中的请求时间，如果没有这个url则返回0

        if time.time() - pended_time < self.pending_threshold: #判断urlpending的时长是否超时，小于阀值就不添加
            print('being downloading:', url)
            return
        if self.db.has(url): ##判断url是否存在数据库中,当has()方法传回True,则进入reture等于不添加，并停止进行下面的代码
            return
        if pended_time: ##经过以上判断，既请求超时从pending中删除，重新加入url池中
            self.pending.pop(url)
        return self.push_to_pool(url)

    def addmany(self, urls, always=False):
        '''判断传入的url是列表还是单个url'''
        if isinstance(urls, str):
            print('urls is a str !!!!', urls)
            self.add(urls, always)
        else:
            for url in urls:
                self.add(url, always)

    def pop(self, count, hub_percent=50):
        '''从url池中获取url给爬取发起请求'''
        print('\n\tmax of host:', self.max_hosts)

        # 取出的url有两种类型：hub=1, 普通=0
        url_attr_url = 0
        url_attr_hub = 1
        # 1. 首先取出hub，保证获取hub里面的最新url.
        hubs = {} #即将抓取的hub的url{hub_uurl:1}
        hub_count = count * hub_percent // 100 #获取一半的hub的url，默认是%50
        for hub in self.hub_pool: #遍历hub的url池 {url: 上一次请求的时间}
            span = time.time() - self.hub_pool[hub]
            if span < self.hub_refresh_span: #判断现在请求的时是否超过请求的间隔时间，如果没有则跳过，继续下一个url
                continue
            hubs[hub] = url_attr_hub  # 1 means hub-url
            self.hub_pool[hub] = time.time() #记录hub-url本次发起请求的时间
            if len(hubs) >= hub_count: #当数量大于1半时停止提取
                break

        # 2. 再取出普通url
        left_count = count - len(hubs) #算出普通url应取出的数量
        urls = {}  #即将抓取的普通url{url:0}
        for host in self.waiting: #取出所有url池的主机号
            if not self.waiting[host]: #如果主机号里没有url，则不再执行以下面的代码，并开始下一组url
                continue
            url = self.waiting[host].pop() #从url池中取出url
            urls[url] = url_attr_url
            self.pending[url] = time.time() ##给本次抓取的url加上现在的时间戳
            if self.max_hosts[0] == host: #实时更新url最多的主机是哪一个
                self.max_hosts[1] -= 1
            if len(urls) >= left_count: #如果取出的数量达到预定量，停止提取
                break
        self.waiting_count -= len(urls) #self.waiting 字典里面的url的个数 减去本次提取的数量
        print('To pop:%s, hubs: %s, urls: %s, hosts:%s' % (count, len(hubs), len(urls), len(self.waiting)))
        urls.update(hubs) #合并两个字典
        return urls

    def size(self,):
        return self.waiting_count

    def empty(self,):
        return self.waiting_count == 0


def test():
    '''测试'''
    pool = UrlPool('crawl_urlpool')
    urls = [
        'http://1.a.cn/xyz',
        'http://2.a.cn/xyz',
        'http://3.a.cn/xyz',
        'http://1.b.cn/xyz-1',
        'http://1.b.cn/xyz-2',
        'http://1.b.cn/xyz-3',
        'http://1.b.cn/xyz-4',
    ]
    pool.addmany(urls)
    del pool

    pool = UrlPool('crawl_urlpool')
    urls = pool.pop(5)
    urls = list(urls.keys())
    print('pop:', urls)
    print('pending:', pool.pending)

    pool.set_status(urls[0], 200)
    print('pending:', pool.pending)
    pool.set_status(urls[1], 404)
    print('pending:', pool.pending)


if __name__ == '__main__':
    test()
